import os.path
from tqdm import tqdm
import shutil, os, json  # 确保所需模块已导入


def traverse_dir(path, tails, file_info_ls=None):
    """
    :param path:
    :param tails:
    :param file_info_ls:
    :return: [(文件路径1，文件名1),(文件路径2，文件名2),(文件路径3，文件名3),...,(文件路径n，文件名n)]
    注意：如路经太深python就无法检测该文件了，遇到这种情况请开启 clean_path功能；
    参考： https://blog.csdn.net/baidu_34928905/article/details/131935428
    """
    if file_info_ls is None:
        file_info_ls = []
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            traverse_dir(file_path, tails, file_info_ls)
        else:
            file_e_ls = os.path.basename(file_path).split('.')
            file_tail = file_e_ls[-1]
            if (tails is None or file_tail.lower() in tails) and not os.path.basename(file_path).startswith('~$'):
                item = (file_path, '.'.join(file_e_ls[0:-1]))
                file_info_ls.append(item)
    return file_info_ls


def parse_img(failed_ls, img_path, out_dir, image_saved_dir):
    # 定义需要处理的图片格式及对应的数据结构
    formats = {
        'jpg': {'file_suffix': 'jpg', 'data': {'file_num': 0, 'file_suffix': 'jpg', 'file_list': []}},
        'jpeg': {'file_suffix': 'jpeg', 'data': {'file_num': 0, 'file_suffix': 'jpeg', 'file_list': []}},
        'png': {'file_suffix': 'png', 'data': {'file_num': 0, 'file_suffix': 'png', 'file_list': []}},
        'bmp': {'file_suffix': 'bmp', 'data': {'file_num': 0, 'file_suffix': 'bmp', 'file_list': []}},
        'jfif': {'file_suffix': 'jfif', 'data': {'file_num': 0, 'file_suffix': 'jfif', 'file_list': []}},
        'exif': {'file_suffix': 'exif', 'data': {'file_num': 0, 'file_suffix': 'exif', 'file_list': []}},
        'gif': {'file_suffix': 'gif', 'data': {'file_num': 0, 'file_suffix': 'gif', 'file_list': []}},
        'tiff': {'file_suffix': 'tiff', 'data': {'file_num': 0, 'file_suffix': 'tiff', 'file_list': []}}
    }

    # 只遍历 formats 中定义的图片后缀
    iter_files = traverse_dir(img_path, tuple(formats.keys()))
    if not iter_files:
        return
    for fid, (file_path, file_uuid) in enumerate(tqdm(iter_files, desc='正在收集图片文件信息')):
        try:
            # 统一采用小写后缀
            img_type = file_path.split('.')[-1].lower()
            if img_type not in formats:
                continue
            # 复制图片到指定目录
            shutil.copyfile(file_path, os.path.join(image_saved_dir, file_uuid))
            image_key = f'image_{img_type}'
            content = f'<image:{image_key}>'
            image_info = {image_key: {'content': '', 'path': file_path}}
            tmp = {
                'file_path': file_path,
                'file_uuid': file_uuid,
                'image_info': image_info,
                'table_info': {},
                'content': content
            }
            formats[img_type]['data']['file_num'] += 1
            formats[img_type]['data']['file_list'].append(tmp)
        except Exception as e:
            failed_ls.append({'filePath': file_path, 'reason': str(e)})

    # 遍历所有图片格式，若有数据则输出到对应的 JSON 文件中
    for ext, info in formats.items():
        if info['data']['file_num'] > 0:
            out_file = os.path.join(out_dir, f'{ext}_content.json')
            with open(out_file, 'w', encoding='utf-8') as f:
                json.dump(info['data'], f, ensure_ascii=False)
